# batch_inference.py
import cv2
import numpy as np
from openvino.runtime import Core
import os
import time

# -----------------------------
# 配置路径（请根据你的实际路径修改！）
# -----------------------------
VEHICLE_MODEL_PATH = r"D:\CodeCNN\intel\vehicle-detection-0202\FP16-INT8\vehicle-detection-0202.xml"
ROAD_MODEL_PATH = r"D:\CodeCNN\intel\road-segmentation-adas-0001\FP16-INT8\road-segmentation-adas-0001.xml"
INPUT_FOLDER = r"H:\xiaomi\test"  # 输入图片文件夹
OUTPUT_FOLDER = r"detect"  # 输出结果文件夹

# 创建输出目录
os.makedirs(OUTPUT_FOLDER, exist_ok=True)

# -----------------------------
# 初始化 OpenVINO
# -----------------------------
core = Core()
device = "CPU"  # 可选 "GPU"（若支持）

# -----------------------------
# 1. 加载车辆检测模型
# -----------------------------
print("Loading vehicle detection model...")
vehicle_model = core.read_model(model=VEHICLE_MODEL_PATH)
vehicle_compiled = core.compile_model(model=vehicle_model, device_name=device)
vehicle_input_layer = vehicle_compiled.input(0)
vehicle_output_layer = vehicle_compiled.output(0)

# 获取输入尺寸
_, _, h_det, w_det = vehicle_input_layer.shape
print(f"Vehicle detection input shape: {h_det}x{w_det}")

# -----------------------------
# 2. 加载道路分割模型
# -----------------------------
print("Loading road segmentation model...")
road_model = core.read_model(model=ROAD_MODEL_PATH)
road_compiled = core.compile_model(model=road_model, device_name=device)
road_input_layer = road_compiled.input(0)
road_output_layer = road_compiled.output(0)

_, _, h_seg, w_seg = road_input_layer.shape
print(f"Road segmentation input shape: {h_seg}x{w_seg}")

# -----------------------------
# 3. 获取 A 文件夹中所有图片
# -----------------------------
image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif'}
image_files = []
for root, dirs, files in os.walk(INPUT_FOLDER):
    for file in files:
        if os.path.splitext(file)[1].lower() in image_extensions:
            image_files.append(os.path.join(root, file))

if not image_files:
    print(f"❌ 没有在 '{INPUT_FOLDER}' 中找到图片！")
    exit(1)

print(f"✅ 发现 {len(image_files)} 张图片，开始批量处理...")

# -----------------------------
# 4. 批量处理每张图片
# -----------------------------
for i, img_path in enumerate(image_files, 1):
    print(f"\n--- 处理第 {i}/{len(image_files)} 张: {os.path.basename(img_path)} ---")

    # 读取图像
    image = cv2.imread(img_path)
    if image is None:
        print(f"⚠️ 无法加载图像: {img_path}")
        continue

    orig_h, orig_w = image.shape[:2]
    print(f"Original size: {orig_w}x{orig_h}")

    # -----------------------------
    # 5. 车辆检测推理
    # -----------------------------
    resized_det = cv2.resize(image, (w_det, h_det))
    input_det = np.expand_dims(resized_det.transpose(2, 0, 1), axis=0)  # HWC -> NCHW
    detections = vehicle_compiled([input_det])[vehicle_output_layer]

    # -----------------------------
    # 6. 道路分割推理
    # -----------------------------
    resized_seg = cv2.resize(image, (w_seg, h_seg))
    input_seg = np.expand_dims(resized_seg.transpose(2, 0, 1), axis=0).astype(np.float32)
    segmentation = road_compiled([input_seg])[road_output_layer]

    seg_mask = segmentation[0, 0]  # HxW
    seg_mask = (seg_mask < 0.5).astype(np.uint8) * 255  # 反转：道路=255，非道路=0
    seg_mask_resized = cv2.resize(seg_mask, (orig_w, orig_h))

    # 创建彩色掩码（绿色道路）
    color_mask = np.zeros_like(image)
    color_mask[:, :, 1] = seg_mask_resized  # Green channel

    # -----------------------------
    # 7. 绘制检测框
    # -----------------------------
    output_image = image.copy()
    for det in detections[0, 0]:
        conf = float(det[2])
        if conf > 0.5:
            x_min = int(det[3] * orig_w)
            y_min = int(det[4] * orig_h)
            x_max = int(det[5] * orig_w)
            y_max = int(det[6] * orig_h)
            cv2.rectangle(output_image, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
            cv2.putText(output_image, f"Car: {conf:.2f}", (x_min, y_min - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)

    # 叠加道路分割（半透明绿色）
    overlay = cv2.addWeighted(output_image, 0.7, color_mask, 0.3, 0)

    # -----------------------------
    # 8. 保存结果到 detect 目录
    # -----------------------------
    base_name = os.path.splitext(os.path.basename(img_path))[0]
    output_filename = f"{base_name}_result.jpg"
    output_path = os.path.join(OUTPUT_FOLDER, output_filename)

    # 如果文件已存在，添加时间戳避免覆盖
    if os.path.exists(output_path):
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        output_filename = f"{base_name}_result_{timestamp}.jpg"
        output_path = os.path.join(OUTPUT_FOLDER, output_filename)

    cv2.imwrite(output_path, overlay)
    print(f"✅ 结果已保存至: {output_path}")

    # -----------------------------
    # 9. （可选）显示最后一张图
    # -----------------------------
    if i == len(image_files):  # 只显示最后一张
        def resize_to_fit(img, target_w=1360, target_h=768):
            h, w = img.shape[:2]
            scale = min(target_w / w, target_h / h)
            new_w, new_h = int(w * scale), int(h * scale)
            resized = cv2.resize(img, (new_w, new_h))
            canvas = np.zeros((target_h, target_w, 3), dtype=np.uint8)
            y_pad = (target_h - new_h) // 2
            x_pad = (target_w - new_w) // 2
            canvas[y_pad:y_pad+new_h, x_pad:x_pad+new_w] = resized
            return canvas

        display_img = resize_to_fit(overlay)
        cv2.namedWindow("Batch Result", cv2.WINDOW_NORMAL)
        cv2.resizeWindow("Batch Result", 1360, 768)
        cv2.imshow("Batch Result", display_img)
        cv2.waitKey(0)
        cv2.destroyAllWindows()

print(f"\n🎉 所有 {len(image_files)} 张图片处理完成！结果保存在 '{OUTPUT_FOLDER}' 目录中。")